package com.qupeng.demo.kafka.kafkaspringbootstreams.wordcount;

import com.qupeng.demo.kafka.kafkaspringbootstreams.common.kafka.streams.JsonSerde;
import com.qupeng.demo.kafka.kafkaspringbootstreams.common.kafka.streams.MyProcessor;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.StreamsBuilder;
import org.apache.kafka.streams.kstream.*;
import org.apache.kafka.streams.state.WindowStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;
import java.util.Arrays;
import java.util.regex.Pattern;

@Configuration
public class WordCountStreams {

    @Bean
    public KStream<String, String> counts(StreamsBuilder myKStreamBuilder) {
        final Pattern pattern = Pattern.compile("\\W+");
        KStream<String, String> source = myKStreamBuilder.stream("words-count-input");
        KStream<String, String> counts = source
                .flatMapValues(value -> Arrays.asList(pattern.split(value.toLowerCase())))
                .map((key, value) -> new KeyValue<String, String>(value, value))
                .filter((key, value) -> !value.equals("the"))
                .groupByKey()
                .count(Named.as("CountStore"))
                .mapValues(value -> Long.toString(value)).toStream();
        counts.print(Printed.toSysOut());
        counts.to("words-count-output");
        return counts;
    }

    @Bean
    public KStream<Windowed<String>, WordsStatistic> windowCounts(StreamsBuilder myKStreamBuilder) {
        final Pattern pattern = Pattern.compile("\\W+");
        KStream<String, String> source = myKStreamBuilder.stream("words-count-input");
        source.process(MyProcessor::new);
        KStream<Windowed<String>, WordsStatistic> counts = source
                .flatMapValues(value -> Arrays.asList(pattern.split(value.toLowerCase())))
                .filter((key, value) -> !value.equals("the"))
                .groupByKey()
                .windowedBy(TimeWindows.ofSizeWithNoGrace(Duration.ofMinutes(5)).advanceBy(Duration.ofSeconds(300)))
                .aggregate(() -> new WordsStatistic(),
                        (key, value, wordsStatistic) -> {
                            wordsStatistic.getWordCount().computeIfPresent(value, (k, v) -> ++v);
                            wordsStatistic.getWordCount().putIfAbsent(value, 1L);
                            return wordsStatistic;
                            },
                        Materialized.<String, WordsStatistic, WindowStore<Bytes, byte[]>>as("count-aggregates").withValueSerde(new JsonSerde(WordsStatistic.class)))
                .toStream();
        counts.print(Printed.toSysOut());
        counts.to("words-count-output", Produced.keySerde(WindowedSerdes.timeWindowedSerdeFrom(String.class, 5*60*1000L)));
        return counts;
    }
}
